from cv2 import cv2
import torch
import torchvision
import torch.optim as optim
from LeNet5 import LeNet5
from utils import one_hot
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# config data
batch_size = 512

# load data
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data/', train=False, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size, shuffle=True)


def showImage(image):
    plt.imshow(image, cmap=plt.cm.gray)
    plt.show()


def model_train():
    # model
    net = LeNet5()
    # Optimizer
    optimizer = optim.Adam(net.parameters(), lr=1e-3)
    for epoch in range(3):
        for batch_idx, (x, y) in enumerate(train_loader):
            out = net(x)
            # [b,10]
            y_onehot = one_hot(y)
            # loss = mse(out,y_onehot)
            loss = net.cretenon(out, y_onehot)
            optimizer.zero_grad()
            loss.backward()
            # w' = w - lr * grad
            optimizer.step()
            if batch_idx % 10 == 0:
                print(epoch, batch_idx, loss.item())
    return net


def model_test():
    model = torch.load('model.tar')
    total_correct = 0
    for x, y in test_loader:
        image = x[0, 0]
        showImage(image.view(28, 28))
        break
        out = model(x)
        pred = out.argmax(dim=1)
        correct = pred.eq(y).sum().float()
        total_correct += correct

    total_sum = len(test_loader.dataset)
    acc = total_correct / total_sum
    print('acc=', acc)


# def loadImage2(path):
    # transform = torchvision.transforms.Compose([
    #     torchvision.transforms.ToTensor(),
    #     torchvision.transforms.Normalize(
    #         (0.1307,), (0.3081,))
    # ])

    # trainset = torchvision.datasets.ImageNet('test.jpg',transform=transform)
    # trans = torchvision.transforms.Compose([
    #     torchvision.transforms.ToTensor(),
    # ])
    # img = trans(gray)
    # img = img.unsqueeze(0)


def laodImage1(path):
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((28, 28)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.1307,), (0.3081,))
    ])
    image = transform(Image.open(path).convert('L'))
    return image.unsqueeze(0)


def laodImage(path):
    image = cv2.imread(path)
    # resize
    image = cv2.resize(image, (28, 28), interpolation=cv2.INTER_LINEAR)
    # 1 channel
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = cv2.adaptiveThreshold(image, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
                                  cv2.THRESH_BINARY, 3, 5)
    image = torch.from_numpy(255-image).unsqueeze(0)
    return image.float().div(255).unsqueeze(0)


def model_deploy():
    model = torch.load('model.tar')
    model.eval()
    image = laodImage('test/6.jpg')
    out = model(image)
    pred = out.argmax(dim=1)
    print(pred.item())


def main():
    # model = model_train()
    # torch.save(model, 'model.tar')
    # model_test()
    model_deploy()


if __name__ == "__main__":
    main()
    cv2.waitKey(0)
